#include "Vec.h"
#include <vector>

class Function{
public:
    virtual ~Function(){}
    virtual Vec<double,6> operator()(const Vec<double,6> &u) const = 0;
};

class JacobiFunction{
public:
    virtual ~JacobiFunction(){}
    virtual Vec<double,36> operator()(const Vec<double,6> &u) const = 0;
};

class F : public Function{
    Vec<double,6> operator()(const Vec<double,6> &u) const override{
        double mu = 1/81.45;
        Vec<double,6> du;
        double d1 = u[1]*u[1]+u[2]*u[2]+(u[0]+mu-1)*(u[0]+mu-1);
        double d2 = u[1]*u[1]+u[2]*u[2]+(u[0]+mu)*(u[0]+mu);
        du[0] = u[3];
        du[1] = u[4];
        du[2] = u[5];
        du[3] = 2*u[4]+u[0]-mu*(u[0]+mu-1)/(d1*sqrt(d1))-(1-mu)*(u[0]+mu)/(d2*sqrt(d2));
        du[4] = -2*u[3]+u[1]-mu*u[1]/(d1*sqrt(d1))-(1-mu)*u[1]/(d2*sqrt(d2));
        du[5] = -mu*u[2]/(d1*sqrt(d1))-(1-mu)*u[2]/(d2*sqrt(d2));
        return du;
    }
};
class JacobiF : public JacobiFunction{
    Vec<double,36> operator()(const Vec<double,6> &u) const override{
        double mu = 1/81.45;
        Vec<double,36> J;
        double d1 = u[1]*u[1]+u[2]*u[2]+(u[0]+mu-1)*(u[0]+mu-1);
        double d2 = u[1]*u[1]+u[2]*u[2]+(u[0]+mu)*(u[0]+mu);
        J[3] = 1-mu/(d1*sqrt(d1))+3*mu*(u[0]+mu-1)*(u[0]+mu-1)/(d1*d1*sqrt(d1))-(1-mu)/(d2*sqrt(d2))+3*(1-mu)*(u[0]+mu)*(u[0]+mu)/(d2*d2*sqrt(d2));
        J[4] = J[9] = 3*mu*u[1]*(u[0]+mu-1)/(d1*d1*sqrt(d1))+3*(1-mu)*u[1]*(u[0]+mu)/(d2*d2*sqrt(d2));
        J[5] = J[15] = 3*mu*u[2]*(u[0]+mu-1)/(d1*d1*sqrt(d1))+3*(1-mu)*u[2]*(u[0]+mu)/(d2*d2*sqrt(d2));
        J[10] = 1-mu/(d1*sqrt(d1))+3*mu*u[1]*u[1]/(d1*d1*sqrt(d1))-(1-mu)/(d2*sqrt(d2))+3*(1-mu)*u[1]*u[1]/(d2*d2*sqrt(d2));
        J[11] = J[16] = 3*mu*u[1]*u[2]/(d1*d1*sqrt(d1))+3*(1-mu)*u[1]*u[2]/(d2*d2*sqrt(d2));
        J[17] = -mu/(d1*sqrt(d1))+3*mu*u[2]*u[2]/(d1*d1*sqrt(d1))-(1-mu)/(d2*sqrt(d2))+3*(1-mu)*u[2]*u[2]/(d2*d2*sqrt(d2));
        J[18] = J[25] = J[32] = 1;
        J[22] = -2;
        J[27] = 2;
        return J;
    }
};
